from openai import OpenAI
from openai.types.chat import ChatCompletion
import os
import json
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

# 自定义配置
BASE_URL: str = os.getenv("OPENAI_BASE_URL", "https://api.chatfire.cn/v1")
API_KEY: str = os.getenv("OPENAI_API_KEY", "sk-zO8exlBicZh7nJeZn5GuC5X9SPuVrZzXoGyOW0i9BFvN62ON")

# 初始化客户端
client: OpenAI = OpenAI(base_url=BASE_URL, api_key=API_KEY)

def generate_atm_requirements():
    prompt = """
请根据以下指令生成一个 ATM 自动取款机系统的需求模型：

1. 包含主要功能模块，例如登录、余额查询、取款、存款、转账、修改密码等。
2. 列出涉及的用户角色，比如普通用户和管理员。
3. 每个功能模块需要有简要的业务流程描述。
4. 使用 JSON 格式输出，不要使用任何 Markdown 或注释。
5. 输出内容应简洁明了，适合后续开发人员理解和实现。

示例输出格式如下：
{
  "system_name": "ATM System",
  "roles": ["Customer", "Admin"],
  "modules": [
    {
      "name": "Login",
      "description": "User authentication using card and PIN.",
      "flow": [
        "Insert card",
        "Enter PIN",
        "Validate credentials",
        "Grant access or deny"
      ]
    }
  ]
}
"""

    try:
        completion: ChatCompletion = client.chat.completions.create(
            model="gpt-4o",
            messages=[
                {"role": "system", "content": "You are a software requirements analyst."},
                {"role": "user", "content": prompt}
            ],
            temperature=0.3,
            max_tokens=1500,
            response_format={"type": "json_object"}
        )

        # 获取响应内容并解析为 JSON
        content = completion.choices[0].message.content.strip()
        requirements = json.loads(content)

        # 打印结果
        print(" 成功生成 ATM 系统需求模型：\n")
        print(json.dumps(requirements, indent=2, ensure_ascii=False))

        # 可选：保存到文件
        with open("atm_requirements.json", "w", encoding="utf-8") as f:
            json.dump(requirements, f, indent=2, ensure_ascii=False)
        print("\n📄 已保存至 atm_requirements.json")

    except Exception as e:
        print(f" 请求失败：{e}")

if __name__ == "__main__":
    generate_atm_requirements()